import numpy as np
import cv2
from PIL import Image, ImageDraw
from .utils import *



class ChestMaskGenerator:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "openpose_json": ("POSE_KEYPOINT",),
                "distance_p": ("FLOAT", {"default": 0.33, "min": 0.0, "max": 1.0, "step": 0.01}),
                "length_p": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01}),
            },
        }

    RETURN_TYPES = ("MASK",)
    FUNCTION = "generate_chest_mask"
    CATEGORY = "🌟FJDH/Mask Generator tool"

    def generate_chest_mask(self, openpose_json,  distance_p, length_p):
        keypoints = openpose_json[0]['people'][0]['pose_keypoints_2d']

        point_1 = np.array(keypoints[3:6])  # Neck
        point_2 = np.array(keypoints[6:9])  # Right shoulder
        point_5 = np.array(keypoints[15:18])  # Left shoulder
        point_8 = np.array(keypoints[24:27])  # Mid-hip

        M = np.linalg.norm(point_1[:2] - point_8[:2])

        N = np.linalg.norm(point_2[:2] - point_5[:2])

        L1_direction = point_5[:2] - point_2[:2]
        L1_direction /= np.linalg.norm(L1_direction)

        if N < 0.2 * M:
            print(f"N is less than 0.2 times M. Setting perpendicular_vector to vertical down.")
            perpendicular_vector = np.array([0, 1])
            L1_direction = np.array([1, 0])  
        else:
            perpendicular_vector = np.array([-L1_direction[1], L1_direction[0]])

        extension_length = M * distance_p
        X = point_1[:2] + perpendicular_vector * extension_length

        Ly_direction = X - point_1[:2]
        Ly_direction /= np.linalg.norm(Ly_direction)

        Ly_length = np.linalg.norm(point_1[:2] - X)

        L2_length = max(N * 2, Ly_length * 2)
        L2_start = X - L1_direction * (L2_length / 2)
        L2_end = X + L1_direction * (L2_length / 2)

        extension_length_Ly = M * length_p
        S_start = L2_start - Ly_direction * extension_length_Ly
        S_end = L2_end - Ly_direction * extension_length_Ly
        S_start_top = L2_start + Ly_direction * extension_length_Ly
        S_end_top = L2_end + Ly_direction * extension_length_Ly

        canvas_width = openpose_json[0]['canvas_width']
        canvas_height = openpose_json[0]['canvas_height']

        color = "white"
        mask = Image.new("RGB", (canvas_width, canvas_height), "black")
        draw = ImageDraw.Draw(mask)

        draw.polygon([tuple(S_start), tuple(S_start_top), tuple(S_end_top), tuple(S_end)], fill=color)

        mask_tensor = pil2tensor(mask)

        mask = mask_tensor[:, :, :, 0]

        return (mask,)

class ForeheadMaskGenerator:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "mask": ("MASK",),
                "points": ("Point",),
                "multiplier": ("INT", {"default": 10, "min": 1, "max": 1000, "step": 1}),
            },
        }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("Mask",)
    FUNCTION = "generate_forehead_mask"
    CATEGORY = "🌟FJDH/Mask Generator tool"

    def generate_forehead_mask(self, mask, points, multiplier):
        all_masks = []

        for batch_mask, batch_points in zip(mask, points):
            height, width = batch_mask.shape[1], batch_mask.shape[2]

            color = "white"
            blank_mask = Image.new("RGB", (width, height), "black")
            draw = ImageDraw.Draw(blank_mask)

            if len(batch_points) == 1:
                mask_np = batch_mask.squeeze().cpu().numpy()
                mask_binary = (mask_np > 0).astype(np.uint8)
                contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
                polygons = [cv2.approxPolyDP(cnt, 0.01 * cv2.arcLength(cnt, True), True) for cnt in contours]
                bboxes = []
                for polygon in polygons:
                    rect = cv2.minAreaRect(polygon)
                    box = cv2.boxPoints(rect)
                    box = np.int0(box)
                    min_x = min(box[:, 0])
                    min_y = min(box[:, 1])
                    max_x = max(box[:, 0])
                    max_y = max(box[:, 1])
                    bboxes.append((min_x, min_y, max_x, max_y))

                bbox_width = bboxes[0][2] - bboxes[0][0]
                bbox_height = bboxes[0][3] - bboxes[0][1]

                x, y = batch_points[0]

                W_y = y - 1 * bbox_height

                rect_height = multiplier * bbox_height

                rect_x1 = 0
                rect_x2 = width
                rect_y1 = W_y - rect_height
                rect_y2 = W_y

                if rect_y1 < 0:
                    rect_y1 = 0

                if rect_y2 < 0:
                    rect_y2 = 0

                draw.rectangle([rect_x1, rect_y1, rect_x2, rect_y2], fill=color)

            elif len(batch_points) >= 2:
                mask_np = batch_mask.squeeze().cpu().numpy()
                mask_binary = (mask_np > 0).astype(np.uint8)
                contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
                polygons = [cv2.approxPolyDP(cnt, 0.01 * cv2.arcLength(cnt, True), True) for cnt in contours]
                bboxes = []
                for polygon in polygons:
                    rect = cv2.minAreaRect(polygon)
                    box = cv2.boxPoints(rect)
                    box = np.int0(box)
                    min_x = min(box[:, 0])
                    min_y = min(box[:, 1])
                    max_x = max(box[:, 0])
                    max_y = max(box[:, 1])
                    bboxes.append((min_x, min_y, max_x, max_y))

                min_bbox_width = min(bbox[2] - bbox[0] for bbox in bboxes)
                min_bbox_height = min(bbox[3] - bbox[1] for bbox in bboxes)

                sorted_points = sorted(batch_points, key=lambda p: p[1])
                A = sorted_points[0]
                B = sorted_points[1]

                A_coords = np.array([A[0], A[1]], dtype=np.float64)
                B_coords = np.array([B[0], B[1]], dtype=np.float64)
                L1_direction = B_coords - A_coords
                L1_length = np.linalg.norm(L1_direction)

                C = (A_coords + B_coords) / 2

                perpendicular_vector = np.array([-L1_direction[1], L1_direction[0]], dtype=np.float64)
                perpendicular_vector /= np.linalg.norm(perpendicular_vector)

                min_value = min(min_bbox_width, min_bbox_height, L1_length)

                if min_value == min_bbox_width or min_value == min_bbox_height:
                    L2_length = 2 * min_value
                else:
                    L2_length = L1_length

                D = C + perpendicular_vector * L2_length

                if D[1] > C[1]:
                    perpendicular_vector = -perpendicular_vector
                    D = C + perpendicular_vector * L2_length

                L3_length = L1_length * 3
                E = D - L1_direction * (L3_length / 2)
                F = D + L1_direction * (L3_length / 2)

                L4_direction_E = E + perpendicular_vector * (multiplier * L1_length)
                L4_direction_F = F + perpendicular_vector * (multiplier * L1_length)

                if L4_direction_E[1] > E[1] or L4_direction_F[1] > F[1]:
                    perpendicular_vector = -perpendicular_vector
                    L4_direction_E = E + perpendicular_vector * (multiplier * L1_length)
                    L4_direction_F = F + perpendicular_vector * (multiplier * L1_length)

                draw.polygon([tuple(E), tuple(F), tuple(L4_direction_F), tuple(L4_direction_E)], fill=color)

            mask_tensor = pil2tensor(blank_mask)
            mask = mask_tensor[:, :, :, 0]
            all_masks.append(mask)

        return (all_masks,)


###########################################################################################




class MaxInscribedRectangleMaskGenerator:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "mask": ("MASK",),
                "width": ("INT", {"default": 100, "min": 1, "max": 10000, "step": 1}),
                "height": ("INT", {"default": 100, "min": 1, "max": 10000, "step": 1}),
                "angle": ("FLOAT", {"default": 0.0, "min": -360.0, "max": 360.0, "step": 0.1}),
                "scale_factor": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
            },
        }

    RETURN_TYPES = ("Point", "Point", "INT", "INT")
    RETURN_NAMES = ("Rectangle-Points", "Rectangle-Centroid", "Rectangle-Width", "Rectangle-Height")
    FUNCTION = "generate_max_inscribed_rectangle_mask"
    CATEGORY = "🌟FJDH/Mask Generator tool"

    def generate_max_inscribed_rectangle_mask(self, mask, width, height, angle, scale_factor):
        all_rectangle_points = []
        all_centroids = []
        all_widths = []
        all_heights = []

        for batch_mask in mask:
            mask_np = batch_mask.squeeze().cpu().numpy()
            mask_binary = (mask_np > 0).astype(np.uint8)

            contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            polygons = [cv2.approxPolyDP(cnt, 0.01 * cv2.arcLength(cnt, True), True) for cnt in contours]

            areas = [cv2.contourArea(polygon) for polygon in polygons]

            if areas:
                max_area_index = np.argmax(areas)
                max_polygon = polygons[max_area_index]
            else:
                all_rectangle_points.append([])
                all_centroids.append([])
                all_widths.append(0)
                all_heights.append(0)
                continue

            aspect_ratio = width / height
            angle_rad = np.deg2rad(angle)
            cos_angle = np.cos(angle_rad)
            sin_angle = np.sin(angle_rad)

            def get_rectangle_corners(center, aspect_ratio, scale, cos_angle, sin_angle):
                half_width = scale * aspect_ratio / 2
                half_height = scale / 2
                corners = np.array([
                    [-half_width, -half_height],
                    [half_width, -half_height],
                    [half_width, half_height],
                    [-half_width, half_height]
                ])
                rotated_corners = np.dot(corners, np.array([[cos_angle, -sin_angle], [sin_angle, cos_angle]]))
                return rotated_corners + center

            def is_inside_polygon(point, polygon):
                return cv2.pointPolygonTest(polygon, tuple(point), False) >= 0

            M = cv2.moments(max_polygon)
            if M["m00"] != 0:
                cX = int(M["m10"] / M["m00"])
                cY = int(M["m01"] / M["m00"])
                center = np.array([cX, cY])

                scale = 1.0
                while True:
                    corners = get_rectangle_corners(center, aspect_ratio, scale, cos_angle, sin_angle)

                    if all(is_inside_polygon(corner, max_polygon) for corner in corners):
                        scale += 0.01
                    else:
                        break

                final_corners = get_rectangle_corners(center, aspect_ratio, scale - 0.01, cos_angle, sin_angle)

                scaled_corners = center + (final_corners - center) * scale_factor
                rectangle_points = [(int(corner[0]), int(corner[1])) for corner in scaled_corners]

                rect_width = int(np.linalg.norm(scaled_corners[0] - scaled_corners[1]))
                rect_height = int(np.linalg.norm(scaled_corners[0] - scaled_corners[3]))

                centroid = [(cX, cY)]

                all_rectangle_points.append(rectangle_points)
                all_centroids.append(centroid)
                all_widths.append(rect_width)
                all_heights.append(rect_height)
            else:
                all_rectangle_points.append([])
                all_centroids.append([])
                all_widths.append(0)
                all_heights.append(0)

        return (all_rectangle_points, all_centroids, all_widths, all_heights)
        


###########################################################################################


class LineMaskGenerator:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "points": ("Point",),
                "line_width": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}),
                "width": ("INT", {"default": 512, "min": 1, "max": 10000, "step": 1}),
                "height": ("INT", {"default": 512, "min": 1, "max": 10000, "step": 1}),
            },
        }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("LineMask",)
    FUNCTION = "generate_line_mask"
    CATEGORY = "🌟FJDH/Mask Generator tool"

    def generate_line_mask(self, points, line_width, width, height):
        masks = []

        for batch_points in points:
            mask = Image.new("RGB", (width, height), "black")
            draw = ImageDraw.Draw(mask)

            for i in range(len(batch_points) - 1):
                start_point = batch_points[i]
                end_point = batch_points[i + 1]
                draw.line([start_point, end_point], fill="white", width=line_width)

            masks.append(mask)

        mask_tensor = pil2tensor(masks)
        mask = mask_tensor[:, :, :, 0]

        return (mask,)

class CircularMaskGenerator:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "points": ("Point",),
                "fill": ("BOOLEAN", {"default": True}),  # New input for fill
                "outline": ("INT", {"default": 0, "min": 0, "max": 10000, "step": 1}),  
                "radius": ("INT", {"default": 5, "min": 1, "max": 10000, "step": 1}),
                "image_width": ("INT", {"default": 512, "min": 1, "max": 10000, "step": 1}),
                "image_height": ("INT", {"default": 512, "min": 1, "max": 10000, "step": 1}),
            }
        }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("Mask",)
    FUNCTION = "generate_circular_mask"
    CATEGORY = "🌟FJDH/Mask Generator tool"

    def generate_circular_mask(self, points, radius, image_width, image_height, fill, outline):
        color = "white" if fill else "black"
        outline_color = "white" if outline > 0 else None

        masks = []

        for batch_points in points:
            mask = Image.new("RGB", (image_width, image_height), "black")
            draw = ImageDraw.Draw(mask)

            for point in batch_points:
                x, y = point
                if outline > 0:
                    draw.ellipse([x - radius, y - radius, x + radius, y + radius], outline=outline_color, width=outline)
                if fill:
                    draw.ellipse([x - radius, y - radius, x + radius, y + radius], fill=color)

            masks.append(mask)

        mask_tensor = pil2tensor(masks)
        mask = mask_tensor[:, :, :, 0]

        return (mask,)

class PolygonMaskGenerator:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "points": ("Point",),
                "image_width": ("INT", {"default": 512, "min": 1, "max": 10000, "step": 1}),
                "image_height": ("INT", {"default": 512, "min": 1, "max": 10000, "step": 1}),
            },
        }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("Mask",)
    FUNCTION = "generate_polygon_mask"
    CATEGORY = "🌟FJDH/Mask Generator tool"

    def generate_polygon_mask(self, points, image_width, image_height):
        all_masks = []

        for batch_points in points:
            color = "white"
            mask = Image.new("RGB", (image_width, image_height), "black")
            draw = ImageDraw.Draw(mask)

            polygon_points = [(point[0], point[1]) for point in batch_points]
            draw.polygon(polygon_points, fill=color)

            all_masks.append(mask)

        mask_tensor = pil2tensor(all_masks)
        mask = mask_tensor[:, :, :, 0]

        return (mask,)

class DistanceMaskGenerator:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "mask": ("MASK",),
                "points": ("Point",),
            },
        }

    RETURN_TYPES = ("MASK",)
    FUNCTION = "generate_distance_mask"
    CATEGORY = "🌟FJDH/Mask Generator tool"

    def generate_distance_mask(self, mask, points):
        all_distance_masks = []

        for batch_mask, batch_points in zip(mask, points):
            mask_np = tensor2np(batch_mask)[0]
            distance_mask = np.zeros_like(mask_np, dtype=np.float32)

            for point in batch_points:
                x, y = point
                distance_map = np.sqrt((np.arange(mask_np.shape[1]) - x)**2 + (np.arange(mask_np.shape[0])[:, np.newaxis] - y)**2)
                distance_mask += distance_map

            distance_mask_min = distance_mask.min()
            distance_mask_max = distance_mask.max()
            distance_mask_normalized = (distance_mask - distance_mask_min) / (distance_mask_max - distance_mask_min)
            distance_mask_tensor = np2tensor(distance_mask_normalized)
            all_distance_masks.append(distance_mask_tensor)

        return (all_distance_masks,)
    



NODE_CLASS_MAPPINGS = {
    "ChestMaskGenerator": ChestMaskGenerator,
    "CircularMaskGenerator": CircularMaskGenerator,
    "ForeheadMaskGenerator": ForeheadMaskGenerator,
    "PolygonMaskGenerator": PolygonMaskGenerator,
    "DistanceMaskGenerator": DistanceMaskGenerator,
    "MaxInscribedRectangleMaskGenerator": MaxInscribedRectangleMaskGenerator,
    "LineMaskGenerator": LineMaskGenerator,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ChestMaskGenerator": "Chest Mask Generator",
    "CircularMaskGenerator": "Circular Mask Generator",
    "ForeheadMaskGenerator": "Forehead Mask Generator",
    "PolygonMaskGenerator": "Polygon Mask Generator",
    "DistanceMaskGenerator": "Distance Mask Generator",
    "MaxInscribedRectangleMaskGenerator": "Max Inscribed Rectangle Mask Generator",
    "LineMaskGenerator": "Line Mask Generator",
}